import json
from collections import deque
from typing import Dict, List, Any, Optional, Callable, Union


class JSONTreeUtil:
    """支持顶层为列表或字典的 JSON 树形结构工具类"""

    @staticmethod
    def from_file(file_path: str, encoding: str = "utf-8") -> Union[Dict[str, Any], List[Any]]:
        """从文件加载数据（顶层可能是 dict 或 list）"""
        with open(file_path, "r", encoding=encoding) as f:
            return json.load(f)

    @staticmethod
    def to_file(data: Union[Dict[str, Any], List[Any]], file_path: str, encoding: str = "utf-8", indent: int = 4) -> None:
        """写入文件"""
        with open(file_path, "w", encoding=encoding) as f:
            json.dump(data, f, indent=indent, ensure_ascii=False)

    @staticmethod
    def _normalize_root(root: Union[Dict[str, Any], List[Any]]) -> List[Dict[str, Any]]:
        """
        统一处理顶层结构：
        - 若顶层是 list，直接返回（视为多个根节点）
        - 若顶层是 dict，包装为 list（视为单个根节点）
        """
        if isinstance(root, list):
            return root
        elif isinstance(root, dict):
            return [root]
        else:
            raise TypeError("顶层结构必须是字典或列表")

    @staticmethod
    def dfs_traverse(
        root: Union[Dict[str, Any], List[Any]],
        callback: Callable[[Dict[str, Any], int], None],
        children_key: str = "children"
    ) -> None:
        """深度优先遍历（支持顶层列表）"""
        roots = JSONTreeUtil._normalize_root(root)

        def _dfs(node: Dict[str, Any], depth: int) -> None:
            callback(node, depth)
            for child in node.get(children_key, []):
                _dfs(child, depth + 1)

        # 遍历所有根节点（顶层列表中的每个元素）
        for r in roots:
            _dfs(r, 0)

    @staticmethod
    def bfs_traverse(
        root: Union[Dict[str, Any], List[Any]],
        callback: Callable[[Dict[str, Any], int], None],
        children_key: str = "children"
    ) -> None:
        """广度优先遍历（支持顶层列表）"""
        roots = JSONTreeUtil._normalize_root(root)
        queue = deque()

        # 初始化队列（所有根节点，深度为 0）
        for r in roots:
            queue.append((r, 0))

        while queue:
            node, depth = queue.popleft()
            callback(node, depth)
            for child in node.get(children_key, []):
                queue.append((child, depth + 1))

    @staticmethod
    def find_node(
        root: Union[Dict[str, Any], List[Any]],
        condition: Callable[[Dict[str, Any]], bool],
        children_key: str = "children"
    ) -> Optional[Dict[str, Any]]:
        """查找满足条件的第一个节点（支持顶层列表）"""
        roots = JSONTreeUtil._normalize_root(root)

        def _search(node: Dict[str, Any]) -> Optional[Dict[str, Any]]:
            if condition(node):
                return node
            for child in node.get(children_key, []):
                result = _search(child)
                if result:
                    return result
            return None

        # 依次搜索每个根节点
        for r in roots:
            found = _search(r)
            if found:
                return found
        return None

    @staticmethod
    def find_nodes(
        root: Union[Dict[str, Any], List[Any]],
        condition: Callable[[Dict[str, Any]], bool],
        children_key: str = "children"
    ) -> List[Dict[str, Any]]:
        """查找所有满足条件的节点（支持顶层列表）"""
        roots = JSONTreeUtil._normalize_root(root)
        all_nodes = []

        def _collect(node: Dict[str, Any]) -> None:
            if condition(node):
                all_nodes.append(node)
            for child in node.get(children_key, []):
                _collect(child)

        for r in roots:
            _collect(r)
        return all_nodes

    @staticmethod
    def add_child(
        root: Union[Dict[str, Any], List[Any]],
        parent_condition: Callable[[Dict[str, Any]], bool],
        new_child: Dict[str, Any],
        children_key: str = "children"
    ) -> bool:
        """向满足条件的父节点添加子节点（支持顶层列表）"""
        parent = JSONTreeUtil.find_node(root, parent_condition, children_key)
        if parent:
            if children_key not in parent:
                parent[children_key] = []
            parent[children_key].append(new_child)
            return True
        return False

    @staticmethod
    def delete_node(
        root: Union[Dict[str, Any], List[Any]],
        condition: Callable[[Dict[str, Any]], bool],
        children_key: str = "children"
    ) -> bool:
        """删除满足条件的节点（支持顶层列表）"""
        roots = JSONTreeUtil._normalize_root(root)

        # 辅助函数：在单个节点及其子树中查找并删除
        def _delete_from_node(current_node: Dict[str, Any]) -> bool:
            # 检查子节点
            children = current_node.get(children_key, [])
            for i, child in enumerate(children):
                if condition(child):
                    del children[i]
                    return True
                # 递归检查子节点的子树
                if _delete_from_node(child):
                    return True
            return False

        # 先检查顶层列表中的根节点
        for i, r in enumerate(roots):
            if condition(r):
                del roots[i]
                return True

        # 再检查每个根节点的子树
        for r in roots:
            if _delete_from_node(r):
                return True

        return False

    @staticmethod
    def to_flat_list(
        root: Union[Dict[str, Any], List[Any]],
        name_key: str = "name",
        children_key: str = "children"
    ) -> List[Dict[str, Any]]:
        """转换为扁平列表（支持顶层列表）"""
        roots = JSONTreeUtil._normalize_root(root)
        flat_list = []

        def _collect(node: Dict[str, Any], parent_path: str) -> None:
            current_name = node[name_key]
            current_path = f"{parent_path}/{current_name}" if parent_path else current_name
            flat_node = node.copy()
            flat_node["path"] = current_path
            flat_list.append(flat_node)
            for child in node.get(children_key, []):
                _collect(child, current_path)

        # 遍历所有根节点，路径前缀为空
        for r in roots:
            _collect(r, "")
        return flat_list

    @staticmethod
    def from_flat_list(
        flat_list: List[Dict[str, Any]],
        name_key: str = "name",
        path_key: str = "path",
        children_key: str = "children"
    ) -> Union[Dict[str, Any], List[Any]]:
        """从扁平列表还原为树形结构（自动判断顶层是 list 还是 dict）"""
        node_map = {}  # 路径 -> 节点
        root_paths = set()  # 根节点路径（无父路径的节点）

        for item in flat_list:
            path = item[path_key]
            node = {k: v for k, v in item.items() if k != path_key}
            node[children_key] = []
            node_map[path] = node

            # 判断是否为根节点（路径中无斜杠）
            if "/" not in path:
                root_paths.add(path)

        # 构建父子关系
        for item in flat_list:
            path = item[path_key]
            if "/" in path:
                parent_path = path.rsplit("/", 1)[0]
                parent_node = node_map.get(parent_path)
                if parent_node:
                    parent_node[children_key].append(node_map[path])

        # 生成结果（若根节点唯一，返回 dict；否则返回 list）
        roots = [node_map[path] for path in root_paths]
        return roots[0] if len(roots) == 1 else roots